import streamlit as st
import asyncio
from llm1 import VolcEngineClient, usellm
from tools2 import generate_echarts_data
import streamlit_echarts
import inspect
import textwrap
from data1 import ST_DATA as ST_DEMOS

# Initialize session state
if 'messages' not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "你好，有什么可以帮助你的？"}]
if 'model_type' not in st.session_state:
    st.session_state.model_type = 'doubao'
if 'long_thinking' not in st.session_state:
    st.session_state.long_thinking = False
if 'chart_type' not in st.session_state:
    st.session_state.chart_type = '柱状图'
if 'function_results' not in st.session_state:
    st.session_state.function_results = []
if 'chart_requested' not in st.session_state:
    st.session_state.chart_requested = False
if 'chart_generated' not in st.session_state:
    st.session_state.chart_generated = False
if 'final_prompt' not in st.session_state:
    st.session_state.final_prompt = ""

# Sidebar
with st.sidebar:
    # Select LLM type
    st.session_state.model_type = st.selectbox(
        "Select LLM type",
        ("doubao", "deepseek"),
        index=0 if st.session_state.model_type == 'doubao' else 1
    )
    # Long thinking button
    st.session_state.long_thinking = st.checkbox("长思考", value=st.session_state.long_thinking)
    # New conversation button
    def clear_chat_history():
        st.session_state.messages = [{"role": "assistant", "content": "你好，有什么可以帮助你的？"}]
        st.session_state.function_results = []
        st.session_state.chart_requested = False
        st.session_state.chart_generated = False
        st.session_state.final_prompt = ""
    st.button("清空聊天记录", on_click=clear_chat_history)
    # Select chart type
    st.session_state.chart_type = st.selectbox(
        "选择 Echarts 图像类型",
        ("柱状图", "折线图", "饼图"),
        index=0 if st.session_state.chart_type == '柱状图' else 1
    )

    st.header("Configuration")

    page_options = list(ST_DEMOS.keys())
    selected_page = st.selectbox(
        label="选择一个问题",
        options=page_options,
        key="select_page"
    )

    selected_province = None
    selected_city = None

    if selected_page == "不同岗位的职位数量":
        selected_province = st.selectbox(
            label="选择一个城市",
            options=["北京", "上海", "天津", "重庆"],
            key="select_province"
        )
    elif selected_page == "不同岗位应聘学历要求分布":
        selected_city = st.selectbox(
            label="选择一个岗位",
            options=["算法工程师", "C/C++开发工程师", "Java工程师", "前端开发工程师", "大数据/数据分析工程师"],
            key="select_city"
        )

    # 侧边栏中的显示图表按钮
    if st.button("显示图表"):
        st.session_state.chart_requested = True
        st.session_state.chart_generated = False

# 执行demo并生成图表
if st.session_state.chart_requested and not st.session_state.chart_generated:
    demo = ST_DEMOS[selected_page][0]

    # 执行 demo 函数并获取结果
    if selected_province:
        result = demo(selected_province)
    elif selected_city:
        result = demo(selected_city)
    else:
        result = demo()

    # 添加图表内容
    chart_content = {"role": "assistant1", "content": f"### {selected_page}", "chart": result}
    st.session_state.messages.append(chart_content)

    st.session_state.chart_generated = True

# 根据状态显示侧边栏提示
if st.session_state.chart_generated:
    st.sidebar.write("当前状态: 图表已生成")
else:
    st.sidebar.write("点击'显示图表'按钮生成分析图表")

# Display chat history
for idx, message in enumerate(st.session_state.messages):
    with st.chat_message(message["role"]):
        st.markdown(message["content"])
        if message["role"] == "assistant1":
            print("实验1数据可视化")
            streamlit_echarts.st_echarts(options=message["chart"])
        # 检查是否有函数调用结果需要绘制图表
        if message["role"] == "assistant" and st.session_state.function_results:
            for result_idx, result in enumerate(st.session_state.function_results):
                if "error" not in result:
                    chart, title = generate_echarts_data(result["result"], st.session_state.chart_type, result["name"])
                    if chart:
                        st.subheader(title)
                        # 添加唯一 key
                        streamlit_echarts.st_pyecharts(chart, key=f"chart_{idx}_{result_idx}")
                    else:
                        st.warning(title)

# Input box
if prompt := st.chat_input("请输入您的问题"):
    # Display user input
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Prepare prompt with history if long thinking is enabled
    if st.session_state.long_thinking:
        history_prompt = "\n".join([f"{msg['role']}: {msg['content']}" for msg in st.session_state.messages])
        final_prompt = f"{history_prompt}\nuser: {prompt}"
    else:
        final_prompt = prompt

    # Save the final prompt to session state
    st.session_state.final_prompt = final_prompt

    # Call LLM to get response
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    results, function_results = loop.run_until_complete(usellm(st.session_state.model_type, final_prompt, st.session_state.chart_type))
    st.session_state.function_results = function_results
    response = results[0] if results else "No valid response was obtained."

    # Display LLM response
    st.session_state.messages.append({"role": "assistant", "content": response})
    with st.chat_message("assistant"):
        st.markdown(response)

        # 绘制图表
        for result in function_results:
            if "error" not in result:
                chart, title = generate_echarts_data(result["result"], st.session_state.chart_type, result["name"])
                if chart:
                    st.subheader(title)
                    streamlit_echarts.st_pyecharts(chart)
                else:
                    st.warning(title)

# Display the final prompt
if st.session_state.final_prompt:
    st.subheader("长思考后发送给AI的问题")
    st.text_area("最终提示信息", st.session_state.final_prompt, height=200)